import os
from dotenv import load_dotenv, find_dotenv
from langchain_openai import OpenAIEmbeddings


_ = load_dotenv(find_dotenv())

api_key = os.environ.get('OPENAI_API_KEY')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")

base_url = os.environ.get('OPENAI_BASE_URL')
if api_key is None:
    raise ValueError("OPENAI_BASE_URL is not set in the .env file")

embeddings  = OpenAIEmbeddings(openai_api_key=api_key, 
                          openai_api_base = base_url,
                          model="text-embedding-3-large")

text = "你变瘦了。"
text_list = [
    "你减肥成功了。",
    "你变苗条了。",
    "你又胖了。",
    "你好瘦啊！"
]

query_result = embeddings.embed_query(text)
# print(f"len(query_result)={len(query_result)}")
# print(f"query_result={query_result}")
doc_result = embeddings.embed_documents(text_list)
# print(f"len(doc_result)={len(doc_result)}")
# print(f"len(doc_result[0])={len(doc_result[0])}")

import numpy as np
def cosine_similarity(vec1, vec2):
    dot_product = np.dot(vec1, vec2)
    norm_vec1 = np.linalg.norm(vec1)
    norm_vec2 = np.linalg.norm(vec2)
    return dot_product / (norm_vec1 * norm_vec2)

for v, t in zip(doc_result, text_list) :
    similarity = cosine_similarity(query_result, v)
    print(f"{t}\t:{similarity}") 
